import TreeModel
import matplotlib.pyplot as plt
import numpy as np


def createDataSet():
    data_set = [[1, 1],
                [1, 1],
                [1, 0],
                [0, 1],
                [0, 1]]
    labels = [1, 1, 0, 0, 0]
    return data_set, labels


if __name__ == '__main__':
    data_set, labels = createDataSet()
    ID3 = TreeModel.ID3()

    print('data_set:\n', data_set)
    print('labels:\n', labels)
    print('ID3 train')
    tree = ID3.train(data_set, labels)

    print('ID3 test')
    print('predict [1,0] ->:', ID3.predict([1, 0]))
    print('predict [1,1] ->:', ID3.predict([1, 1]))

    print('CART tree')
    cart = TreeModel.CART()
    train_set = np.linspace(0, 2, 200)
    train_labels = np.sin(6 * train_set) + 0.2 * np.random.random(train_set.shape)

    perm = np.random.permutation(200)
    train_set = train_set.reshape(200, 1)
    train_set = train_set[perm]
    train_labels = train_labels[perm]

    train_set, test_set = train_set[:150], train_set[150:]
    train_labels, test_labels = train_labels[:150], train_labels[150:]

    # 蓝色点表示测试集
    plt.plot(test_set, test_labels, 'b.')
    # 红色点表示训练集
    plt.plot(train_set, train_labels, 'r.')

    print('CART tree is training')
    cart.train(train_set, train_labels)

    t_set = np.linspace(0, 2, 200).reshape(200, 1)

    p_set = []
    for x in t_set:
        p_set.append(cart.predict(x))

    # 黄色折线表示未剪枝的回归树
    plt.plot(t_set, p_set, 'y-')

    cart.tree = cart.prune(cart.tree, test_set, test_labels)
    p_set = []
    for x in t_set:
        p_set.append(cart.predict(x))
    # 绿色折线代表剪枝后的回归树
    plt.plot(t_set, p_set, 'g-')

    plt.show()

